package com.fangdd.logtrace.jms;

import com.fangdd.logtrace.BoundsLog;
import com.fangdd.logtrace.TraceId;
import com.fangdd.logtrace.TraceIdThreadLocal;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * Created by sulong on 16-9-12.
 */
public class TraceIdJmsSenderWrapper {
    private static final Logger LOGGER = LoggerFactory.getLogger(TraceIdJmsSenderWrapper.class);

    public static void wrap(final String name, final CallWith consumer) {
        final TraceId traceId = TraceIdThreadLocal.get();

        TraceId next = null;
        if (!traceId.isUnset()) {
            next = TraceIdThreadLocal.sequentialNext();
        }

        long startTime = System.currentTimeMillis();
        try {
            if (!traceId.isUnset()) {
                BoundsLog.callerIn(LOGGER, next.toString(), name);
            }

            consumer.with(next);

            if (!traceId.isUnset()) {
                BoundsLog.callerOut(LOGGER, next.toString(), name, true, System.currentTimeMillis() - startTime);
            }
        } catch (RuntimeException runtimeException) {
            if (!traceId.isUnset()) {
                BoundsLog.callerOut(LOGGER, next.toString(), name, false, System.currentTimeMillis() - startTime);
            }
            throw runtimeException;
        }
    }

    public interface CallWith {
        public void with(TraceId traceId);
    }
}
